from abc import ABC

register_map = {
    'qwen': 'QwenAgent',
    'qa_outline_story_writer': 'QAOutlineStoryWriter',
    'musicgen_t2m': 'MusicGenAgent',
    'story_diffusion_t2i': 'StoryDiffusionAgent',
    'cosyvoice_tts': 'CosyVoiceAgent',
    'audioldm2_t2a': 'AudioLDM2Agent',
    'slideshow_video_compose': 'SlideshowVideoComposeAgent',
    'freesound_sfx_retrieval': 'FreesoundSfxAgent',
    'freesound_music_retrieval': 'FreesoundMusicAgent',
}    


def import_from_register(key):
    #value = register_map[key]
    # exec(f'from . import {value}')
    value = register_map[key]
    try:
        # 尝试直接导入
        exec(f'from . import {value}')
    except ImportError as e:
        print(f'第一次导入 {key} 失败: {e}')
        try:
            # 如果失败，尝试从具体模块导入
            if key == 'slideshow_video_compose':
                exec('from .video_compose_agent import SlideshowVideoComposeAgent')
            elif key == 'story_diffusion_t2i':
                exec('from .image_generation_agent import StoryDiffusionAgent')
            # 添加其他工具的特殊处理...
        except ImportError as e2:
            print(f'导入 {key} 完全失败: {e2}')
            raise e2


class ToolRegistry(dict):

    def _import_key(self, key):
        try:
            import_from_register(key)
        except Exception as e:
            print(f'import {key} failed, details: {e}')

    def __getitem__(self, key):
        if key not in self.keys():
            self._import_key(key)
        return super().__getitem__(key)

    def __contains__(self, key):
        self._import_key(key)
        return super().__contains__(key)
    

TOOL_REGISTRY = ToolRegistry()


def register_tool(name):

    def decorator(cls):
        TOOL_REGISTRY[name] = cls
        return cls
    
    return decorator


def init_tool_instance(cfg):
    return TOOL_REGISTRY[cfg["tool"]](cfg["cfg"])